/**
 * @file base64.c
 * @author bobwxc (bobwxc@yeah.net)
 * @brief base64 encode
 * @version 0.1
 * @date 2022-11-15
 *
 * @copyright Copyright (c) 2022 bobwxc@yeah.net, Public Domain
 *
 */

#include "base64.h"

const char base64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/**
 * @brief base64 3 byte process core
 *
 * @param src 3 byte data, input
 * @param res 4 byte char, output
 * @return int =0 success
 */
int b64_3(unsigned char *src, unsigned char *res)
{
	unsigned char a, b, c, d;
	a = (src[0] >> 2) & 0b00111111;

	b = (src[0] << 4) & 0b00111111;
	b += (src[1] >> 4) & 0b00001111;

	c = (src[1] << 2) & 0b00111111;
	c += (src[2] >> 6) & 0b00000011;

	d = src[2] & 0b00111111;

	res[0] = base64_table[a];
	res[1] = base64_table[b];
	res[2] = base64_table[c];
	res[3] = base64_table[d];

	return 0;
}

/**
 * @brief base64 encode
 *
 * @param src
 * @param len length of src
 * @param res output, include terminal zero
 * @param max_len max length of res
 * @return char* res, =NULL for error
 */
char *base64_encode(unsigned char *src, size_t len, char *res, size_t max_len)
{
	size_t return_len = len / 3 * 4;
	short rest_len = len % 3;

	if ((rest_len == 0 && max_len <= return_len) || (rest_len != 0 && max_len <= (return_len + 4)))
	{
		res[0] = '\0';
		return NULL;
	}

	for (size_t i = 0; i < (len / 3); i++)
		b64_3(src + (i * 3), res + (i * 4));

	if (rest_len == 1)
	{
		res[return_len] = base64_table[(src[len - rest_len] >> 2) & 0b00111111];
		res[return_len + 1] = base64_table[(src[len - rest_len] << 4) & 0b00111111];
		res[return_len + 2] = '=';
		res[return_len + 3] = '=';
		res[return_len + 4] = '\0';
	}
	else if (rest_len == 2)
	{
		res[return_len] = base64_table[(src[len - rest_len] >> 2) & 0b00111111];
		res[return_len + 1] = base64_table[((src[len - rest_len] << 4) & 0b00111111) + ((src[len - rest_len + 1] >> 4) & 0b00001111)];
		res[return_len + 2] = base64_table[(src[len - rest_len + 1] << 2) & 0b00111111];
		res[return_len + 3] = '=';
		res[return_len + 4] = '\0';
	}

	return res;
}

const unsigned char base64_suffix_map[] = {
	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	62, // '+'
	0, 0, 0,
	63,										// '/'
	52, 53, 54, 55, 56, 57, 58, 59, 60, 61, // '0'-'9'
	0, 0, 0, 0, 0, 0, 0,
	0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, // 'A'-'Z'
	0, 0, 0, 0, 0, 0,
	26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, // 'a'-'z'
};

/**
 * @brief base64 decode
 *
 * @param src do NOT check data
 * @param len length of src, exclude terminal zero
 * @param res
 * @param max_len max length of res
 * @return int actual length of res
 */
size_t base64_decode(char *src, size_t len, unsigned char *res, size_t max_len)
{
	if (len % 4 != 0 || max_len < (len / 4 * 3))
	{
		res[0] = '\0';
		return -1;
	}

	size_t i = 0;

	for (i = 0; i < (len / 4 - 1); i++)
	{
		if (src[i * 4] > 122 || src[i * 4 + 1] > 122 || src[i * 4 + 2] > 122 || src[i * 4 + 3] > 122)
		{
			res[0] = '\0';
			return -1;
		}
		res[i * 3] = (base64_suffix_map[(int)src[i * 4]] << 2) + ((base64_suffix_map[(int)src[i * 4 + 1]] >> 4) & 0b00000011);
		res[i * 3 + 1] = (base64_suffix_map[(int)src[i * 4 + 1]] << 4) + ((base64_suffix_map[(int)src[i * 4 + 2]] >> 2) & 0b00001111);
		res[i * 3 + 2] = (base64_suffix_map[(int)src[i * 4 + 2]] << 6) + base64_suffix_map[(int)src[i * 4 + 3]];
	}

	i = (len / 4 - 1);

	if (src[i * 4 + 2] == '=' && src[i * 4 + 3] == '=')
	{
		if (src[i * 4] > 122 || src[i * 4 + 1] > 122)
		{
			res[0] = '\0';
			return -1;
		}
		res[i * 3] = (base64_suffix_map[(int)src[i * 4]] << 2) + ((base64_suffix_map[(int)src[i * 4 + 1]] >> 4) & 0b00000011);
		i = i * 3 + 1;
	}
	else if (src[i * 4 + 3] == '=')
	{
		if (src[i * 4] > 122 || src[i * 4 + 1] > 122 || src[i * 4 + 2] > 122)
		{
			res[0] = '\0';
			return -1;
		}
		res[i * 3] = (base64_suffix_map[(int)src[i * 4]] << 2) + ((base64_suffix_map[(int)src[i * 4 + 1]] >> 4) & 0b00000011);
		res[i * 3 + 1] = (base64_suffix_map[(int)src[i * 4 + 1]] << 4) + ((base64_suffix_map[(int)src[i * 4 + 2]] >> 2) & 0b00001111);
		i = i * 3 + 2;
	}
	else
	{
		if (src[i * 4] > 122 || src[i * 4 + 1] > 122 || src[i * 4 + 2] > 122 || src[i * 4 + 3] > 122)
		{
			res[0] = '\0';
			return -1;
		}
		res[i * 3] = (base64_suffix_map[(int)src[i * 4]] << 2) + ((base64_suffix_map[(int)src[i * 4 + 1]] >> 4) & 0b00000011);
		res[i * 3 + 1] = (base64_suffix_map[(int)src[i * 4 + 1]] << 4) + ((base64_suffix_map[(int)src[i * 4 + 2]] >> 2) & 0b00001111);
		res[i * 3 + 2] = (base64_suffix_map[(int)src[i * 4 + 2]] << 6) + base64_suffix_map[(int)src[i * 4 + 3]];
		i = i * 3 + 3;
	}

	return i;
}
